# 2025-02-18
# 主要是一些基本的,常用的逻辑处理的函数,它们不要引用以下其他特殊的库

import numpy as np
import os
import pandas as pd
import random
import glob
# 2025-02-18


def headers_similar(header1, header2, allowed_diff=1):
    """
    判断两个列表 header1 和 header2 是否“相似”:
    要求二者长度一致,并且对应位置上不同的元素个数不超过 allowed_diff 个.
    """
    if len(header1) != len(header2):
        return False
    diff_count = sum(1 for a, b in zip(header1, header2) if a != b)
    return diff_count <= allowed_diff
# 2025-02-18


def collect_target_dirs(root_folder):
    """
    遍历 root_folder 下的所有子文件夹,查找每个子文件夹中 analysis 子目录下存在的 txt 文件,
    并读取其文件头(列名).这里不依赖于文件夹名称,而是根据“文件字段是否基本一致”
    来判断是否属于同一批目标文件.只有当同一批(即字段相似)的文件数量 >= 2 时,
    才认为它们是目标,函数返回这些文件所在的文件夹完整路径列表.

    参数:
      - root_folder: 包含各个待处理文件夹的根目录

    返回:
      - target_dirs: 符合条件(含有目标文件)的文件夹完整路径列表
    """
    # 存储候选文件夹及其对应的文件字段(只取 analysis 目录下第一个能读取到 header 的 txt 文件)
    candidate_dirs = []  # 元素形如 (folder_path, header_list)

    # 遍历根目录下的所有子文件夹(只考虑一级目录)
    for subfolder in sorted(os.listdir(root_folder)):
        subfolder_path = os.path.join(root_folder, subfolder)
        if os.path.isdir(subfolder_path):
            analysis_path = os.path.join(subfolder_path, "analysis")
            if os.path.isdir(analysis_path):
                found_candidate = False
                # 遍历 analysis 目录中的文件
                for file in sorted(os.listdir(analysis_path)):
                    if file.lower().endswith('.txt'):
                        file_path = os.path.join(analysis_path, file)
                        try:
                            # 仅读取表头(假设文件为制表符分隔)
                            df = pd.read_csv(
                                file_path, sep='\t', engine='python', nrows=0)
                            header = list(df.columns)
                            candidate_dirs.append((subfolder_path, header))
                            found_candidate = True
                            break  # 每个文件夹只取第一个能成功读取 header 的 txt 文件
                        except Exception as e:
                            print(f"读取文件 {file_path} 表头时出错: {e}")
                if not found_candidate:
                    # 若 analysis 文件夹内没有符合条件的 txt 文件,则跳过该文件夹
                    # print(f"文件夹 {analysis_path} 内未找到合适的 txt 文件.")
                    pass

    # 根据文件头将候选文件夹分组,相似(字段差异不超过 1 个)的归为一组
    groups = []  # 每组为列表,元素为 (folder_path, header)
    for folder_path, header in candidate_dirs:
        placed = False
        for group in groups:
            # 以组中第一个文件的 header 作为代表
            rep_header = group[0][1]
            if headers_similar(header, rep_header, allowed_diff=1):
                group.append((folder_path, header))
                placed = True
                break
        if not placed:
            groups.append([(folder_path, header)])

    # 只保留那些组内文件数量大于等于 2 的文件夹
    target_dirs = []
    for group in groups:
        if len(group) >= 2:
            for folder_path, _ in group:
                target_dirs.append(folder_path)

    return target_dirs
# 2025-02-21


def extract_timesteps(xyz_path):
    # 2025-02-21
    # 读取 xyz 文件的 ITEM: TIMESTEP 时间信息,速度快
    timesteps = []

    with open(xyz_path, 'r') as file:
        lines = file.readlines()
        i = 0  # 当前行号

        while i < len(lines):
            # 查找 TIMESTEP 行
            if lines[i].startswith('ITEM: TIMESTEP'):
                # 下一行就是实际的 TIMESTEP 数值
                timestep = int(lines[i+1].strip())
                timesteps.append(timestep)

                # 跳过 ITEM: NUMBER OF ATOMS 行及后面的 9 行
                i += 3  # 跳过 'ITEM: TIMESTEP' 和对应的数值行
                # 读取 ITEM: NUMBER OF ATOMS
                number_of_atoms = int(lines[i].strip())
                i += 9  # 跳过接下来的 9 行 (包括箱体边界和原子数据)
            else:
                i += 1  # 如果不是 TIMESTEP 行,继续往下读

    return timesteps
# 2025-02-23 查找 批任务 中的各文件


def find_lammps_files(space_path):
    """
    遍历指定目录,直到找到所有的 input.lammps 文件.
    参数:
        space_path (str): 起始的目录路径.
    返回:
        list: 包含所有 input.lammps 文件完整路径的列表(已排序).
    """
    # 初始化一个列表来存储所有找到的 input.lammps 文件路径
    lammps_paths = []

    # 遍历 space_path 目录及其子目录
    for root, dirs, files in os.walk(space_path):
        for file in files:
            # 如果文件名是 input.lammps
            if file == "input.lammps":
                # 记录文件的完整路径
                lammps_paths.append(root)

    # 对路径列表进行排序
    lammps_paths.sort()

    return lammps_paths

# 2025-03-17 判断一维离散轨迹从首点到尾点是否在同一侧.


def trajectory_endpoints_side(points, L):
    """
    判断一维离散轨迹从首点到尾点是否在同一侧.

    参数:
      points: 在区间 (0, L) 内的离散点列表(顺序给出轨迹)
      L: PCB 的长度,左边界为 0,右边界为 L.

    算法:
      对于每一对相邻点,计算 d = x[i+1] - x[i].
        - 如果 d > L/2,则说明走的直接距离太远,实际最短路径为 d - l,此时认为发生了一次向左的跨界(计数 -1).
        - 如果 d < -L/2,则实际最短路径为 d + l,此时计数 +1.
        - 否则,不发生跨界.

      累加所有跨界计数,若总跨界次数为偶数,则首尾点在同一侧;奇数则在不同侧.
    """
    crossing_count = 0
    for i in range(len(points) - 1):
        d = points[i+1] - points[i]
        if d > L/2:
            # 例如从 0.1 到 0.9:d=0.8 > 0.5,最短路径为 0.8 - L = -0.2,认为跨界一次(向左跨界)
            crossing_count -= 1
        elif d < -L/2:
            # 例如从 0.9 到 0.1:d=-0.8 < -0.5,最短路径为 -0.8 + L = 0.2,认为跨界一次(向右跨界)
            crossing_count += 1
    if crossing_count % 2 == 0:
        return "same"
    else:
        return "different"


# 2025-04-14 角度格式判断
def get_box_dimensions(box):
    """
    判断 box 格式并返回盒子的尺寸:

    - 如果 box 是坐标格式(前 3 个数字为起始坐标,后 3 个数字为结束坐标),
      则返回每个轴上结束坐标减去起始坐标的差值.

    - 如果 box 是长度+角度格式(前 3 个数字代表长度,后 3 个数字代表角度):
      首先判断后三个数字是否均为 10 的倍数:
        - 如果满足,返回前 3 个数字(长度).
        - 如果不满足,返回“数值异常”.

    参数:
      box: array-like,包含 6 个数(如 float32)

    返回:
      如果 box 为坐标格式,返回 [box[3]-box[0], box[4]-box[1], box[5]-box[2]];
      如果为长度+角度格式且角度均为 10 的倍数,则返回前 3 个数字;
      否则返回字符串 "数值异常".
    """
    box = np.array(box, dtype=np.float32)

    if box.shape[0] != 6:
        raise ValueError("box 必须有 6 个元素!")

    # 判断后三个数字是否都是 10 的倍数
    angles = box[3:6]
    if np.all(angles % 10 == 0):
        return box[0:3]
    # 判断是否为坐标格式(起始坐标应均小于结束坐标)
    elif box[0] <= box[3] and box[1] <= box[4] and box[2] <= box[5]:
        dims = box[3:6] - box[0:3]
        return dims
    else:
        return "数值异常"

# 2025-04-29 判断位置属性的名称格式


def detect_position_key(graph):
    for node, data in graph.nodes(data=True):
        if "position" in data:
            return "position"
        elif "posi" in data:
            return "posi"
        elif "pos" in data:
            return "pos"
    return None  # 如果都没有找到

# 2025-05-07 计算网络的应变率


def compute_box_strain_rate(initial_graph, current_graph) -> float:
    """
    计算两个网络图在 x 方向盒子尺寸的应变率.

    参数:
        initial_graph (nx.MultiGraph): 初始态网络,要求 initial_graph.graph['box'] 存放盒子信息.
        current_graph (nx.MultiGraph): 当前态网络,要求 current_graph.graph['box'] 存放盒子信息.

    返回:
        float: x 方向的应变率,(L_current - L_initial) / L_initial.
    """
    # 从 metadata 中读取盒子尺寸(返回 [Lx, Ly, Lz])
    L_initial = get_box_dimensions(initial_graph.graph['box'])[0]
    L_current = get_box_dimensions(current_graph.graph['box'])[0]

    return (L_current - L_initial) / L_initial


def sample_missing(
    missing,
    threshold: int = 500,
    max_samples: int = 2000,
    seed: int = None
):
    """
    如果 missing 的长度 <= threshold，直接返回原序列（保持顺序）；
    否则，随机抽样 min(len(missing), max_samples) 个元素返回（无放回）。

    参数
    ----
    missing : 序列
        待抽样的元素容器（如 list, tuple, np.ndarray 等）。
    threshold : int
        开始抽样的阈值（默认 500）。
    max_samples : int
        抽样上限，最终返回的元素数不会超过此值（默认 2000）。
    seed : int | None
        随机种子，方便复现。

    返回
    ----
    List[T]
        抽样结果列表，长度 ≤ max_samples。
    """
    # 如果是集合类型，先转换为列表
    if isinstance(missing, set):
        missing = list(missing)

    N = len(missing)
    # 小数量时不抽样
    if N <= threshold:
        return list(missing)

    # 确定抽样数目
    M = min(N, max_samples)
    rnd = random.Random(seed)

    # 从 [0..N-1] 中无放回抽取 M 个索引
    idx = rnd.sample(range(N), M)
    # 按抽样索引取元素
    if N > threshold:
        print(f"抽样比例: {M/N:.2%}")
    return [missing[i] for i in idx]

# 2025-06-14 获得不同链接率下的初始 graph 字典


def inital_graph_paths(base_dir):
    """
    初始化并获取所有graph文件路径

    Args:
        base_dir (str): 基础目录路径

    Returns:
        dict: 包含交联度和对应graph文件路径的字典
    """
    # 获取所有bondingprob目录
    bonding_dirs = glob.glob(os.path.join(base_dir, "string=bondingprob*/N0"))

    # 存储graph文件路径的字典
    graph_files = {}
    for dir_path in sorted(bonding_dirs):
        bonding_prob = dir_path.split("bondingprob")[1].split("/")[0]
        graph_path = os.path.join(dir_path, "graph_series", "0-graph.pkl")
        if os.path.exists(graph_path):
            graph_files[bonding_prob] = graph_path

    return graph_files
